﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using Mono.Cecil;
using Mono.Cecil.Cil;
using NETDeob.Deobfuscators;
using NETDeob.Engine.Utils;
using NETDeob.Misc;
using NETDeob.Misc.Structs__Enums___Interfaces.Deobfuscation;

namespace NETDeob.Core.Deobfuscators.Manco.Tasks
{
    class StringDecryptor : DeobfuscationTask
    {
        private class MancoEntry
        {
            public Algorithm Algo;
            
            public string Key;
            public string OldString;
        }

        public enum Algorithm
        {
            Plain = 0,
            DES = 1,
            RC2 = 2,
            Rijndael = 4,
            TripleDES = 8
        }

        public StringDecryptor(AssemblyDefinition asmDef)
            : base(asmDef)
        {
            RoutineDescription = "String decryption";
        }

        [DeobfuscationPhase(1, "Locating decryptor method")]
        public bool Phase1()
        {
            string key;
            Algorithm strAlgo;

            var decryptor = FindDecryptorMethod(out key, out strAlgo);
            PhaseParam = new object[] {decryptor, key, strAlgo};

            if(decryptor == null){
                ThrowPhaseError("Could not locate decryptor method!", 1, false);
                return false;
            }

            MarkMember(decryptor);

            if(key == null){
                ThrowPhaseError("Key is null!", 1, false);
                return false;
            }

            return true;
        }

        [DeobfuscationPhase(2, "Decrypting & marking strings")]
        public bool Phase2()
        {
            var decryptor = PhaseParam[0];
            var key = PhaseParam[1];
            var strAlgo = PhaseParam[2];

            var decryptedFields = DecryptStrings(decryptor.DeclaringType, key, strAlgo);

            PhaseParam = new object[] {decryptedFields, PhaseParam[2]};

            Logger.VSLog(string.Format("{0} strings were decrypted and marked for replacement...", decryptedFields.Count));
            return true;
        }

        [DeobfuscationPhase(3, "Replacing encrypted strings")]
        public bool Phase3()
        {
            ReplaceStrings();

            Logger.VSLog("All marked strings are replaced with plaintext...");
            return true;
        }

        [DeobfuscationPhase(4, "Cleaning up", true)]
        public bool CleanUp()
        {
            var badType = PhaseParam;
            badType.Module.Types.Remove(badType);

            Logger.VSLog(string.Format("Cleaned {0} bad methods, {1} bad fields & 1 bad type from the assembly...",
                                       badType.Methods.Count, badType.Fields.Count));

            return true;
        }

        #region Reversed methods

        private static string DecryptString(string paramStr, string key, Algorithm strAlgo)
        {
            #region Plain

            if (strAlgo == Algorithm.Plain)
            if (paramStr != string.Empty)
            {
                int num2;
                int length = (paramStr.Length/2) + (paramStr.Length%2);
                string str = paramStr.Substring(0, length);
                string str2 = paramStr.Substring(length);
                StringBuilder builder = new StringBuilder(paramStr.Length);
                char[] chArray = new char[str.Length];
                char[] chArray2 = new char[str2.Length];
                bool flag = (str.Length%2) == 0;
                for (num2 = 0; num2 < str.Length; num2 += 2)
                {
                    chArray[num2/2] = str[num2];
                    if ((num2 < (str.Length - 1)) || flag)
                    {
                        chArray[(str.Length - (num2/2)) - 1] = str[num2 + 1];
                    }
                }
                flag = (str2.Length%2) == 0;
                for (num2 = 0; num2 < str2.Length; num2 += 2)
                {
                    chArray2[num2/2] = str2[num2];
                    if ((num2 < (str2.Length - 1)) || flag)
                    {
                        chArray2[(str2.Length - (num2/2)) - 1] = str2[num2 + 1];
                    }
                }
                flag = str.Length == str2.Length;
                for (num2 = 0; num2 < str.Length; num2++)
                {
                    builder.Append(chArray[num2]);
                    if ((num2 < (str.Length - 1)) || flag)
                    {
                        builder.Append(chArray2[num2]);
                    }
                }
                return builder.ToString();
            }

            #endregion

            if (strAlgo != Algorithm.Plain)
            {
                var transform = CreateAlgo(key, strAlgo).CreateDecryptor();
                var buf = Convert.FromBase64String(paramStr);

                return Encoding.Default.GetString(transform.TransformFinalBlock(buf, 0, buf.Length));
            }

            ThrowPhaseError("Failed to decrypt string!", 0, false);

            return "undecrypted";
        }
     
        #endregion

        private static void ReplaceStrings()
        {
            var decryptedFields = PhaseParam[0];

            foreach (var mDef in from modDef in AsmDef.Modules from typeDef in modDef.Types from mDef in typeDef.Methods where mDef.HasBody select mDef)
                for (int i = 0; i < mDef.Body.Instructions.Count; i++)
                {
                    var instr = mDef.Body.Instructions[i];

                    if (instr.OpCode != OpCodes.Ldsfld) continue;
                    if (instr.Operand.ToString().ToLower().Contains("string::empty"))
                        continue;

                    if (!decryptedFields.ContainsKey(instr.Operand as FieldDefinition)) continue;
                    var ilProc = mDef.Body.GetILProcessor();

                    ilProc.InsertBefore(instr,
                                        ilProc.Create(OpCodes.Ldstr,
                                                      decryptedFields[
                                                          (PhaseParam = (instr.Operand as FieldDefinition))]));
                    MarkMember(instr, mDef);
                }

            PhaseParam = PhaseParam.DeclaringType;
        }
        private static SymmetricAlgorithm CreateAlgo(string paramStr, Algorithm strAlgo)
        {
            SymmetricAlgorithm algorithm = GetAlgo(strAlgo);
            ASCIIEncoding encoding = new ASCIIEncoding();
            algorithm.IV = encoding.GetBytes(paramStr.Substring(0, (strAlgo == Algorithm.Rijndael ? 0x10 : 8)));
            algorithm.Key = encoding.GetBytes(paramStr.Substring((strAlgo == Algorithm.Rijndael ? 0x10 : 8)));
            return algorithm;
        }
        private static SymmetricAlgorithm GetAlgo(Algorithm strAlgo)
        {
            switch (strAlgo)
            {
                case Algorithm.DES:
                    return new DESCryptoServiceProvider();

                case  Algorithm.RC2:
                    return new RC2CryptoServiceProvider();

                case Algorithm.Rijndael:
                    return new RijndaelManaged();

                case Algorithm.TripleDES:
                    return new TripleDESCryptoServiceProvider();
            }

            return null;
        }

        public static Dictionary<FieldDefinition, string> DecryptStrings(TypeDefinition type, string key, Algorithm strAlgo)
        {
            var outDict = new Dictionary<FieldDefinition, string>();
            var curInstr = type.Methods.First(method => method.Name == ".cctor").Body.Instructions[0];

            while(curInstr.Next.OpCode != OpCodes.Ret)
            {
                if(curInstr.OpCode == OpCodes.Ldstr)
                {
                    var decryptedStr = "";
                  
                    switch(strAlgo)
                    {
                        case Algorithm.Plain:
                            decryptedStr = DecryptString(curInstr.Operand as string, "", strAlgo);
                            outDict.Add(curInstr.Next.Next.Operand as FieldDefinition, decryptedStr);
                        break;

                        case Algorithm.RC2:
                        case Algorithm.DES:
                        case Algorithm.Rijndael:
                        case Algorithm.TripleDES:

                            decryptedStr = DecryptString(curInstr.Operand as string, key, strAlgo);
                            outDict.Add(curInstr.Next.Next.Operand as FieldDefinition, decryptedStr);

                        break;
                    }

                    Logger.VLog(string.Format(@"[Decrypt] ""{0}"" -> ""{1}""", curInstr.Operand as string, decryptedStr));
                }

                curInstr = curInstr.Next;
            }

            return outDict;
        }

        public static MethodDefinition FindDecryptorMethod(out string key, out Algorithm strAlgo)
        {
            strAlgo = Algorithm.Plain;

            var nonCipherSignature = new ILSignature
                                {
                                    StartIndex = 1,
                                    StartOpCode = OpCodes.Nop,
                                    Instructions = new List<OpCode>
                                                       {
                                                           OpCodes.Ldarg_0,
                                                           OpCodes.Ldsfld,
                                                           OpCodes.Call,
                                                           OpCodes.Ldc_I4_0,
                                                           OpCodes.Ceq,
                                                           OpCodes.Stloc_S,
                                                           OpCodes.Ldloc_S,
                                                           OpCodes.Brtrue,
                                                           OpCodes.Nop
                                                       }
                                };

            foreach (var modDef in AsmDef.Modules)
                foreach (var typeDef in modDef.Types)
                    foreach (var mDef in typeDef.Methods)
                        if (SignatureFinder.IsMatch(mDef, nonCipherSignature))
                        {
                            Logger.VSLog("Found decryptor module at " + mDef.Name + "...");
                            Logger.VSLog("Strings are plainly encoded...");
                            key = "";

                            return mDef;
                        }

            var cipherSignature = new ILSignature
                                      {
                                          StartIndex = 2,
                                          StartOpCode = OpCodes.Nop,
                                          Instructions = new List<OpCode>
                                                             {
                                                                 OpCodes.Ldarg_0,
                                                                 OpCodes.Ldstr,
                                                                 OpCodes.Call,
                                                                 OpCodes.Call,
                                                                 OpCodes.Stloc_0
                                                             }
                                      };

            foreach (var modDef in AsmDef.Modules)
                foreach (var typeDef in modDef.Types)
                    foreach (var mDef in typeDef.Methods)
                        if (SignatureFinder.IsMatch(mDef, cipherSignature))
                        {
                            var encType = "";
                            Logger.VSLog("Found decryptor module at " + mDef.Name + "...");

                            if (mDef.Body.Instructions[1].OpCode == OpCodes.Ldc_I4_0)
                            {
                                encType = "DES";
                                strAlgo = Algorithm.DES;
                            }
                            else if (mDef.Body.Instructions[1].OpCode == OpCodes.Ldc_I4_1)
                            {
                                encType = "RC2";
                                strAlgo = Algorithm.RC2;
                            }
                            else if (mDef.Body.Instructions[1].OpCode == OpCodes.Ldc_I4_2)
                            {
                                encType = "Rijndael";
                                strAlgo = Algorithm.Rijndael;
                            }
                            else if (mDef.Body.Instructions[1].OpCode == OpCodes.Ldc_I4_3)
                            {
                                encType = "TripleDES";
                                strAlgo = Algorithm.TripleDES;
                            }

                            key = mDef.Body.Instructions.FirstOfOpCode(OpCodes.Ldstr).Operand as string;

                            Logger.VSLog("Strings are encrypted using a " + encType + " algorithm...");

                            return mDef;
                        }
            key = "";
            return null;
        }

        public void InitializeDecryption(object param)
        {
            
        }

        public void DecryptEntry(ref object entry)
        {
            throw new NotImplementedException();
        }

        public void ProcessEntry(object entry)
        {
            throw new NotImplementedException();
        }

    }
}
